import ml_collections
import torch
import torchvision
import torch.optim as optim
import torch.nn as nn
from torchvision import datasets
from matplotlib import pyplot as plt
from dataset import getDataLoader
from resnet import ResNet18
from FcaNet.fcanet import fcanet18
from coordatt import CoordNet18
from ViT import VisionTransformer
from torchinfo import summary
import time
import os

MODEL_PATH = 'saved_models'
def save_model(model, type, name='model', path=MODEL_PATH):
    if not os.path.exists(f'{path}/{type}'):
        os.makedirs(f'{path}/{type}')
    # 获取当前时间
    t = time.localtime()
    model_path = f'{path}/{type}/{t.tm_year}-{t.tm_mon}-{t.tm_mday}-{t.tm_hour}-{t.tm_min}-{t.tm_sec}'
    os.makedirs(model_path)
    # 将模型保存到指定路径
    torch.save(model.state_dict(), model_path + f'/{name}.pth')


def train_one_epoch(epoch, model, trainloader, criterion, optimizer, losses, times, device, name='cnn'):
    model.train()
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        # get the inputs; data is a list of [inputs, labels]
        inputs, labels = data[0].to(device), data[1].to(device)

        # zero the parameter gradients
        optimizer.zero_grad()

        # forward + backward + optimize
        outputs = model(inputs)
        if name == 'cnn':
            loss = criterion(outputs, labels)
        else:
            loss = criterion(outputs[0], labels)
        loss.backward()
        optimizer.step()

        # print statistics
        running_loss += loss.item()
        if i % times == times-1:
            print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / times:.3f}')
            losses.append(running_loss / times)
            running_loss = 0.0

def eval_model(model, testloader, accuracys1, accuracys5, topk=(1,5), device='cpu', name='cnn'):
    model.eval()
    correct_1 = 0
    correct_5 = 0
    total = 0
    with torch.no_grad():
        for data in testloader:
            images, labels = data[0].to(device), data[1].to(device)
            outputs = model(images)
            if name == 'cnn':
                _, pred = outputs.topk(max(topk), 1, True, True)
            else:
                _, pred = outputs[0].topk(max(topk), 1, True, True)
            pred = pred.t()
            correct = pred.eq(labels.view(1, -1).expand_as(pred))

            for k in topk:
                correct_k = correct[:k].reshape(-1).float().sum(0, keepdim=True).item()
                if k == 1:
                    correct_1 += correct_k
                elif k == 5:
                    correct_5 += correct_k
            total += labels.size(0)

    print(f'Top-1 accuracy of the network on the 10000 test images: {100 * correct_1 / total:.2f} %')
    print(f'Top-5 accuracy of the network on the 10000 test images: {100 * correct_5 / total:.2f} %')
    accuracys1.append(100 * correct_1 / total)
    accuracys5.append(100 * correct_5 / total)


def plot_loss(losses,accuracys1,accuracys5):
    plt.figure(figsize=(15,5))

    plt.subplot(1, 3, 1)
    plt.plot(losses)
    plt.xlabel("every 128*125 item")
    plt.ylabel("loss")
    plt.grid()

    plt.subplot(1, 3, 2)
    plt.plot(accuracys1)
    plt.xlabel("every epoch")
    plt.ylabel("Top-1 accuracy")
    plt.grid()

    plt.subplot(1, 3, 3)
    plt.plot(accuracys5)
    plt.xlabel("every epoch")
    plt.ylabel("Top-5 accuracy")
    plt.grid()

def train(num_epoch, model, trainloader, testloader, criterion, optimizer, times, device, name='cnn'):
    losses = []
    accuracys1 = []
    accuracys5 = []
    for epoch in range(num_epoch):
        train_one_epoch(epoch, model, trainloader, criterion, optimizer, losses, times, device, name)
        eval_model(model, testloader, accuracys1, accuracys5, device=device, name=name)
    plot_loss(losses, accuracys1, accuracys5)


# train the mini ViT
config = ml_collections.ConfigDict()
config.patches = ml_collections.ConfigDict({'size': (4, 4)})
config.hidden_size = 252
config.transformer = ml_collections.ConfigDict()
config.transformer.mlp_dim = 1536
config.transformer.num_heads = 4
config.transformer.num_layers = 4
config.transformer.attention_dropout_rate = 0.0
config.transformer.dropout_rate = 0.1
config.classifier = 'token'
config.representation_size = None


if __name__ == "__main__":

    # learningj_rate = 1e-3
    learning_rate = 2e-4
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # 修改下面一行代码，使其调用你的模型
    model = CoordNet18(100)
    model = model.to(device)
    name = "cnn"

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

    trainLoader, testLoader = getDataLoader()

    num_epoch = 10
    times = 16000 // 128
    train(num_epoch, model, trainLoader, testLoader, criterion, optimizer, times, device, name)
    name = 'coord'
    save_model(model, name)